package com.xiashitech.agent.instrumentation.utils;

import com.xiashitech.agent.instrumentation.base.context.TraceConst;
import com.xiashitech.agent.instrumentation.base.context.TraceContextVars;
//import com.xiashitech.agent.instrumentation.base.context.bridge.BridgeBaggage;
//import com.xiashitech.agent.instrumentation.base.context.bridge.BridgeBaggageBuilder;
//import com.xiashitech.agent.instrumentation.bridge.Java8Bridge;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.api.baggage.Baggage;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.api.baggage.BaggageBuilder;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.api.baggage.BaggageEntry;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.Context;

import java.util.Iterator;
import java.util.Map;

public class BaggageUtil {

    /**
     * 获取当前的baggage，自动合并TraceContextVars中未同步的变量
     */
    public static Baggage getCurrentBaggage() {
        Baggage baggage = Baggage.current();

        TraceContextVars contextVars = TraceContextVars.get();
        if (contextVars.size() == 0)
            return baggage;

        String vStr = baggage.getEntryValue(TraceConst.XS_CONTEXT_VARS_VERSION);
        if (vStr == null || vStr.isEmpty() || Integer.parseInt(vStr) < contextVars.getVersion()) {
            BaggageBuilder builder = baggage.toBuilder();
            contextVars.forEach((n, v) -> { //TODO REENTRY变量是否需要放入baggage?
                if(needsSync(n.getName(), v.getValue()))
                    builder.put(n.getName(), v.getValue());
            });
            builder.put(TraceConst.XS_CONTEXT_VARS_VERSION, String.valueOf(contextVars.getVersion()));
            baggage = builder.build();
        }

        return baggage;
    }

    /**
     * 兼容代码，从TraceContextVars中同步数据
     */
    public static BaggageBuilder syncBaggage() {
        Baggage baggage = Baggage.current();
        BaggageBuilder baggageBuilder = baggage.toBuilder();

        TraceContextVars contextVars = TraceContextVars.get();
        if (contextVars.size() == 0)
            return baggageBuilder;

        String vStr = baggage.getEntryValue(TraceConst.XS_CONTEXT_VARS_VERSION);
        if (vStr == null || vStr.isEmpty() || Integer.parseInt(vStr) < contextVars.getVersion()) {
            contextVars.forEach((n, v) -> { //TODO REENTRY变量是否需要放入baggage?
                if(needsSync(n.getName(), v.getValue()))
                    baggageBuilder.put(n.getName(), v.getValue());
            });
            baggageBuilder.put(TraceConst.XS_CONTEXT_VARS_VERSION, String.valueOf(contextVars.getVersion()));
        }

        return baggageBuilder;
    }

    private static boolean needsSync(String name, String value) {
        if("endpoint-api".equals(name))
            return false;
        if (value == null || value.isEmpty())
            return false;
        return true;
    }

    /**
     * 获取Trace变量的值，从TraceContextVars和Baggage中查找
     */
    public static String getTraceVar(Context context, String name) {
        String value = TraceContextVars.get().get(name);
        if(value == null || value.isEmpty())
            value = Baggage.fromContext(context).getEntryValue(name);
        return value;
    }

//    /**
//     * 将BridgeBaggage同步到Baggage
//     *
//     * @return
//     */
//    public static BaggageBuilder syncBaggage() {
//        BaggageBuilder baggageBuilder = Java8Bridge.currentForBaggage().toBuilder();
//        if (Java8Bridge.currentForBridgeBaggage().size() > 0) {
//            Map<String, BaggageEntry> mapBridgeBaggage = Java8Bridge.currentForBridgeBaggage().asMap();
//            Iterator<String> iteratorBridgeBaggage = Java8Bridge.currentForBridgeBaggage().asMap().keySet().iterator();
//            while (iteratorBridgeBaggage.hasNext()) {
//                String key = iteratorBridgeBaggage.next();
//                if("endpoint-api".equals(key))
//                    continue;
//                BaggageEntry value = mapBridgeBaggage.get(key);
//                if (value == null || value.getValue().isEmpty())
//                    continue;
//                baggageBuilder.put(key, value.getValue());
//            }
//        }
//        return baggageBuilder;
//    }

//    /**
//     * 获取BridgeBaggageBuilder，用以构建BridgeBaggage
//     *
//     * @return
//     */
//    public static BridgeBaggageBuilder getBridgeBaggage() {
//        return BridgeBaggage.current().toBuilder();
//    }

    /**
     * 从指定Context提取Baggage，然后合并到现有Baggage中
     *
     * @param baggage
     * @param context
     * @return
     */
    public static BaggageBuilder mergeBaggageFromContext(Baggage baggage, Context context) {
        if (baggage == null) {
            return null;
        }
        BaggageBuilder baggageBuilder = baggage.toBuilder();
        Baggage.fromContext(context).forEach((key, value) -> baggageBuilder.put(key, value.getValue()));
        return baggageBuilder;
    }
}
